#include "tensorflow_plugin/src/profiler/cpu/utils/xplane_builder.h"

#include <algorithm>
#include <string>
#include <utility>

#include "absl/container/flat_hash_map.h"
#include "absl/strings/string_view.h"
#include "tensorflow_plugin/src/utils/types.h"
#include "tensorflow_plugin/src/profiler/cpu/utils/time_utils.h"
#include "tensorflow_plugin/src/utils/xplane.pb.h"
namespace demo_plugin {
namespace profiler {

XPlaneBuilder::XPlaneBuilder(XPlane* plane)
    : XStatsBuilder<XPlane>(plane, this), plane_(plane) {
  for (auto& id_and_metadata : *plane->mutable_event_metadata()) {
    auto& metadata = id_and_metadata.second;
    last_event_metadata_id_ =
        std::max<int64>(last_event_metadata_id_, metadata.id());
    if (!metadata.name().empty()) {
      event_metadata_by_name_.try_emplace(metadata.name(), &metadata);
    }
  }
  for (auto& id_and_metadata : *plane->mutable_stat_metadata()) {
    auto& metadata = id_and_metadata.second;
    last_stat_metadata_id_ =
        std::max<int64>(last_stat_metadata_id_, metadata.id());
    if (!metadata.name().empty()) {
      stat_metadata_by_name_.try_emplace(metadata.name(), &metadata);
    }
  }
  for (XLine& line : *plane->mutable_lines()) {
    lines_by_id_.try_emplace(line.id(), &line);
  }
}

XEventMetadata* XPlaneBuilder::GetOrCreateEventMetadata(int64 metadata_id) {
  XEventMetadata& metadata = (*plane_->mutable_event_metadata())[metadata_id];
  metadata.set_id(metadata_id);
  return &metadata;
}

XEventMetadata* XPlaneBuilder::CreateEventMetadata() {
  return GetOrCreateEventMetadata(++last_event_metadata_id_);
}

XEventMetadata* XPlaneBuilder::GetOrCreateEventMetadata(
    absl::string_view name) {
  XEventMetadata*& metadata = event_metadata_by_name_[name];
  if (metadata == nullptr) {
    metadata = CreateEventMetadata();
    metadata->set_name(std::string(name));
  }
  return metadata;
}

XEventMetadata* XPlaneBuilder::GetOrCreateEventMetadata(std::string&& name) {
  XEventMetadata*& metadata = event_metadata_by_name_[name];
  if (metadata == nullptr) {
    metadata = CreateEventMetadata();
    metadata->set_name(std::move(name));
  }
  return metadata;
}

XStatMetadata* XPlaneBuilder::GetOrCreateStatMetadata(int64 metadata_id) {
  XStatMetadata& metadata = (*plane_->mutable_stat_metadata())[metadata_id];
  metadata.set_id(metadata_id);
  return &metadata;
}

XStatMetadata* XPlaneBuilder::CreateStatMetadata() {
  return GetOrCreateStatMetadata(++last_stat_metadata_id_);
}

XStatMetadata* XPlaneBuilder::GetOrCreateStatMetadata(absl::string_view name) {
  XStatMetadata*& metadata = stat_metadata_by_name_[name];
  if (metadata == nullptr) {
    metadata = CreateStatMetadata();
    metadata->set_name(std::string(name));
  }
  return metadata;
}

XStatMetadata* XPlaneBuilder::GetOrCreateStatMetadata(std::string&& name) {
  XStatMetadata*& metadata = stat_metadata_by_name_[name];
  if (metadata == nullptr) {
    metadata = CreateStatMetadata();
    metadata->set_name(std::move(name));
  }
  return metadata;
}

XLineBuilder XPlaneBuilder::GetOrCreateLine(int64 line_id) {
  XLine*& line = lines_by_id_[line_id];
  if (line == nullptr) {
    line = plane_->add_lines();
    line->set_id(line_id);
  }
  return XLineBuilder(line, this);
}

XEventBuilder XLineBuilder::AddEvent(const XEventMetadata& metadata) {
  XEvent* event = line_->add_events();
  event->set_metadata_id(metadata.id());
  return XEventBuilder(line_, plane_, event);
}

XEventBuilder XLineBuilder::AddEvent(const XEvent& event) {
  XEvent* new_event = line_->add_events();
  *new_event = event;
  return XEventBuilder(line_, plane_, new_event);
}

void XLineBuilder::SetTimestampNsAndAdjustEventOffsets(int64 timestamp_ns) {
  int64 offset_ps = NanosToPicos(line_->timestamp_ns() - timestamp_ns);
  line_->set_timestamp_ns(timestamp_ns);
  if (offset_ps) {
    for (auto& event : *line_->mutable_events()) {
      event.set_offset_ps(event.offset_ps() + offset_ps);
    }
  }
}

}  // namespace profiler
}  // namespace demo_plugin
